"""
ConnLLM Moonshot流式输出能力V2实现
提供Moonshot API流式输出功能
"""

from typing import Dict, Any, Generator, List, Optional
import json
import logging

from ...core_v2.capability import ModelCapability

logger = logging.getLogger("connllm.implementations.moonshot.streaming_v2")

class MoonshotStreamingV2(ModelCapability):
	"""Moonshot流式输出能力V2实现"""
	
	def stream_response(self, messages: List[Dict[str, Any]]) -> Generator[str, None, None]:
		"""
		流式输出响应
		
		Args:
			messages: 消息列表
			
		Returns:
			文本片段生成器
		"""
		# 准备流式请求参数
		params = self.provider._prepare_request_params(messages)
		params["stream"] = True
		
		# 获取API端点
		endpoint = self.provider._get_completion_endpoint()
		
		try:
			# 创建请求上下文
			context = {"original_messages": messages}
			
			# 应用请求处理器
			for processor in self.provider.request_processors:
				params = processor.process_request(params, context)
			
			# 发送流式请求
			accumulated_text = ""
			current_chunk = ""
			chunks_count = 0
			
			for chunk in self.provider.http_client.stream("POST", endpoint, params):
				try:
					# 解析数据块
					chunk_str = chunk.decode("utf-8").strip()
					if not chunk_str:
						continue
						
					# 处理SSE格式
					if chunk_str.startswith("data: "):
						chunk_str = chunk_str[6:]
						
					# 处理[DONE]标记
					if chunk_str == "[DONE]":
						break
						
					# 解析JSON数据
					data = json.loads(chunk_str)
					
					# 提取文本
					delta_text = self._extract_text_from_delta(data)
					if delta_text:
						accumulated_text += delta_text
						current_chunk += delta_text
						chunks_count += 1
						
						# 每累积一定量的文本或达到一定数量的块后输出
						if len(current_chunk) > 10 or chunks_count % 5 == 0:
							yield current_chunk
							current_chunk = ""
				except Exception as e:
					logger.warning(f"处理流式块失败: {str(e)}")
			
			# 输出剩余的文本
			if current_chunk:
				yield current_chunk
				
			logger.debug(f"流式输出完成，共{len(accumulated_text)}字符")
			
		except Exception as e:
			logger.error(f"流式输出失败: {str(e)}")
			# 返回错误信息
			yield f"[流式输出错误: {str(e)}]"
	
	def _extract_text_from_delta(self, delta_data: Dict[str, Any]) -> str:
		"""
		从delta数据中提取文本
		
		Args:
			delta_data: delta数据
			
		Returns:
			提取的文本
		"""
		try:
			if "choices" in delta_data and len(delta_data["choices"]) > 0:
				choice = delta_data["choices"][0]
				
				# 检查delta格式
				if "delta" in choice and "content" in choice["delta"]:
					content = choice["delta"]["content"]
					if content is not None:
						return content
		except Exception as e:
			logger.error(f"从delta中提取文本失败: {str(e)}")
			
		return ""
